package com.wxc;

import com.wxc.printer.BinaryTreeInfo;

import java.util.Comparator;

/**
 * @author wxc
 * @date 2021年08月28日 19:47
 * 大顶堆
 */
public class BinaryHeap<E> extends AbstractHeap<E> implements BinaryTreeInfo<Integer> {
    private final int DEFAULT_CAPACITY = 10;
    protected E[] elements;
    // 默认大顶堆
    private boolean flag = true;

    public BinaryHeap(Comparator<E> comparator, E[] elements) {
        this(comparator, elements, elements.length,true);
    }

    // 用于TOP K只复制前面一些元素
    public BinaryHeap(Comparator<E> comparator, E[] elements, int copySize,boolean flag) {
        super(comparator);
        this.flag = flag;
        if (elements == null || elements.length == 0) this.elements = (E[]) new Object[DEFAULT_CAPACITY];
        else {
            this.size = copySize;
            int capacity = copySize < DEFAULT_CAPACITY ? DEFAULT_CAPACITY : copySize + (copySize << 1);
            this.elements = (E[]) new Object[capacity];
            System.arraycopy(elements, 0, this.elements, 0, copySize);
            heapify();
        }
    }

    public BinaryHeap() {
        this(null, null);
    }

    public void setFlag(boolean flag) {
        this.flag = flag;
    }

    private int compare(E e1, E e2) {
        if(flag) {
            if (comparator != null) return comparator.compare(e1, e2);
            if (e1 instanceof Comparable) return ((Comparable<E>) e1).compareTo(e2);
        }
        else {
            if (comparator != null) return -comparator.compare(e1, e2);
            if (e1 instanceof Comparable) return -((Comparable<E>) e1).compareTo(e2);
        }
        throw new UnsupportedOperationException();
    }

    private void emptyCheck() {
        if (size == 0) throw new IndexOutOfBoundsException();
    }

    private void elementNotNullCheck(E element) {
        if (element == null) throw new IllegalArgumentException();
    }

    private void ensureCapacity(int capacity) {
        int oldCapacity = elements.length;
        if (capacity > oldCapacity) {
            E[] newElements = (E[]) new Object[oldCapacity + (oldCapacity >> 1)];
            if (size >= 0) System.arraycopy(elements, 0, newElements, 0, size);
            elements = newElements;
        }
    }

    // 批量建堆
    private void heapify() {
        // 自上而下的上滤
//        for (int i = 1; i < size; i++) {
//            siftUp(i);
//        }
        // 自下而上的下滤
        for (int i = (size >> 1) - 1; i >= 0; i--) {
            siftDown(i);
        }
    }

    // 上滤
    private void siftUp(int index) {
        E element = elements[index];
        while (index > 0) {
            int parentIndex = (index - 1) >> 1;
            E parentElement = elements[parentIndex];
            if (compare(element, parentElement) < 0) break;
            // swap
            elements[index] = parentElement;
            elements[parentIndex] = element;
            index = parentIndex;
        }
    }

    // 优化 swap
    // 2 * logn -> 1 * logn + 1
    private void siftUpPlus(int index) {
        E element = elements[index];
        while (index > 0) {
            int parentIndex = (index - 1) >> 1;
            E parentElement = elements[parentIndex];
            // 如果插入结点小于父节点 直接退出
            // 我这里如果相等那就覆盖
            if (compare(element, parentElement) < 0) break;
            // swap
            elements[index] = parentElement;
            index = parentIndex;
        }
        elements[index] = element;
    }

    private void siftDown(int index) {
        E element = elements[index];
        // index 要小于第一个叶子结点的索引 <==> 小于所有非叶子节点的数量 <==> floor(size / 2)
        while (index < (size >> 1)) {
            // 1.只有左子结点
            // 2.左右子结点都有
            int leftIndex = (index << 1) + 1;
            E leftNode = elements[leftIndex];
            int maxIndex = leftIndex;
            E maxNode = leftNode;

            int rightIndex = leftIndex + 1;
            // 右边比左边大
            if (rightIndex < size) {
                E rightNode = elements[rightIndex];
                if (compare(rightNode, leftNode) > 0) {
                    maxIndex = rightIndex;
                    maxNode = rightNode;
                }
            }
            // 比左右都大(不可能 因为最后一个结点肯定是左右某个结点的子孙节点)
            // 只有相等的可能
            if (compare(element, maxNode) > 0) break;
            // 将子节点存放到index位置
            elements[index] = maxNode;
            // 重新设置index
            index = maxIndex;
        }
        elements[index] = element;
    }


    @Override
    public void add(E element) {
        elementNotNullCheck(element);
        ensureCapacity(size + 1);
        elements[size++] = element;
        siftUpPlus(size - 1);
    }

    @Override
    public E get() {
        emptyCheck();
        return elements[0];
    }

    @Override
    public E remove() {
        emptyCheck();
        E root = elements[0];
        // 最后一个进行覆盖
        int lastIndex = --size;
        elements[0] = elements[lastIndex];
        elements[lastIndex] = null;
        siftDown(0);
        return root;
    }

    @Override
    public E replace(E element) {
        elementNotNullCheck(element);
        // 这样确实可以 但是慢 2 * logn
        // E root = remove();
        // add(element);

        // 直接将堆顶元素变成新添加元素 然后下溢
        // 1 * logn;
        E root = size == 0 ? null : elements[0];
        elements[0] = element;
        siftDown(0);
        return root;
    }

    @Override
    public void clear() {
        for (int i = 0; i < size; i++) elements[i] = null;
        size = 0;
    }

    @Override
    public Integer root() {
        return 0;
    }

    @Override
    public Integer left(Integer index) {
        int leftIndex = (index << 1) + 1;
        return leftIndex >= size ? null : leftIndex;
    }

    @Override
    public Integer right(Integer index) {
        int leftIndex = (index << 1) + 2;
        return leftIndex >= size ? null : leftIndex;
    }

    @Override
    public String string(Integer index) {
        return elements[index].toString();
    }
}
